package pool

import (
	"context"
	"log"
	"sync"
	"sync/atomic"
	"time"
)

type TaskFn func(ctx context.Context)

type WorkerPool struct {
	minWorkerNum int32
	maxWorkerNum int32
	poolSize     int32
	inWorkNum    int32
	step         int32
	taskPulisher chan TaskFn
	taskReciever chan TaskFn
	isStop       bool
	taskQueue    []TaskFn
	ctx          context.Context
	cancel       context.CancelFunc
	taskLock     sync.Mutex
}

func NewWorkerPool(min, max, step int32) *WorkerPool {
	ctx, cancel := context.WithCancel(context.Background())
	return &WorkerPool{
		maxWorkerNum: max,
		minWorkerNum: min,
		poolSize:     0,
		inWorkNum:    0,
		taskPulisher: make(chan TaskFn, 1),
		taskReciever: make(chan TaskFn, 1),
		ctx:          ctx,
		cancel:       cancel,
		step:         step,
		taskQueue:    make([]TaskFn, 0),
		taskLock:     sync.Mutex{},
	}
}

func (wp *WorkerPool) Stop() {
	wp.isStop = true
}
func (wp *WorkerPool) Pulish(task TaskFn) {
	if !wp.isStop {
		wp.taskReciever <- task
	}
}
func (wp *WorkerPool) init() {
	wp.createWorker(wp.minWorkerNum)
}

func (wp *WorkerPool) Poster() {
	for !wp.isStop {
		select {
		case task := <-wp.taskReciever:
			wp.taskLock.Lock()
			wp.taskQueue = append(wp.taskQueue, task)
			wp.taskLock.Unlock()
		case <-wp.ctx.Done():

		}
	}
}

func (wp *WorkerPool) createWorker(num int32) {
	for i := 0; i < int(num); i++ {
		go func(ctx context.Context) {
			var isStop bool
			for !isStop {
				select {
				case <-ctx.Done():
					isStop = true
				case fn := <-wp.taskPulisher:
					if fn == nil {
						isStop = true
						return
					}
					atomic.AddInt32(&wp.inWorkNum, 1)
					func() {
						defer func() {
							msg := recover()
							log.Default().Println(msg)
						}()
						fn(ctx)
					}()
					atomic.AddInt32(&wp.inWorkNum, -1)
				}
			}
		}(wp.ctx)
		wp.poolSize++
	}
}

func (wp *WorkerPool) Run() {
	wp.init()
	go wp.Poster()

	for {
		//check
		if len(wp.taskQueue) >= int(wp.step) {
			if wp.inWorkNum == wp.poolSize {
				newSize := wp.poolSize + wp.step
				if newSize > wp.maxWorkerNum {
					newSize = wp.maxWorkerNum
				}
				wp.createWorker(newSize - wp.poolSize)
			}
		}

		if len(wp.taskQueue) > 0 {
			wp.taskLock.Lock()
			task := wp.taskQueue[0]
			wp.taskQueue = wp.taskQueue[1:]
			wp.taskLock.Unlock()
			wp.taskPulisher <- task
		} else {
			if wp.isStop {
				wp.cancel()
			} else {
				time.Sleep(time.Millisecond * 10)
			}
		}

	}
}
